#!/usr/bin/python

# This module is used to query interface and cpu data via pysnmp

from pysnmp.entity.rfc3413.oneliner import cmdgen


class QueryOidData:

    def __init__(self, ip, community):
        self.ip = ip
        self.community = community
        self.ifnum = self.query_ifnum()
        self.ifnames = self.query_ifname()

    # The cmdgenerator used to query data per oid
    def cmdgenerator(self, oid):
        generator = cmdgen.CommandGenerator()
        comm_data = cmdgen.CommunityData('server', self.community, 1)
        transport = cmdgen.UdpTransportTarget((self.ip, 161))
        real_fun = getattr(generator, 'getCmd')
        errorindication, errorstatus, errorindex, varbinds = real_fun(comm_data, transport, oid)
        if errorindication is not None or errorstatus is True or varbinds[0][1].prettyPrint() == "No Such Instance currently exists at this OID":
            return None
        else:
            return varbinds[0][1].prettyPrint()

    # Fetch the number of interfaces
    def query_ifnum(self):
        oid = (1, 3, 6, 1, 2, 1, 2, 1, 0)
        interf_num = self.cmdgenerator(oid)
        return interf_num

    # Fetch interface names
    def query_ifname(self):
        interf_names = []
        if self.ifnum is not None:
            oids = [(1, 3, 6, 1, 2, 1, 31, 1, 1, 1, 1, i) for i in range(1, (int(self.ifnum)+1))]
            interf_names = filter(lambda x: x is not None, map(lambda oid: self.cmdgenerator(oid), oids))
        return interf_names

    # Fetch inbound and outbound traffic of interfaces
    def query_iftraffic(self):
        """
        Data format: {'eth0':{
                              'in':'123', 
                              'out':'321'
                              },
                      'eth1':{
                             'in':'456', 
                             'out':'654'
                             }
                      }        
        """
        if_traffic = {}
        if self.ifnum is not None:
            for i in range(1, (int(self.ifnum)+1)):
                interface_name = self.ifnames[i-1]
                if_traffic[interface_name] = {}
                ifin_oid = (1, 3, 6, 1, 2, 1, 2, 2, 1, 10, i)
                ifout_oid = (1, 3, 6, 1, 2, 1, 2, 2, 1, 16, i)
                ifin_traff = self.cmdgenerator(ifin_oid)
                ifout_traff = self.cmdgenerator(ifout_oid)
                if_traffic[interface_name]['in'] = ifin_traff
                if_traffic[interface_name]['out'] = ifout_traff
        return if_traffic

    # Fetch cpu utilization
    def query_cpuutil(self):
        cpu_oid = (1, 3, 6, 1, 4, 1, 2021, 11, 10, 0)
        cpu_uti = self.cmdgenerator(cpu_oid)
        return cpu_uti

